#include "bzbmem/mem.h"
#include "bzbfile/file.h"
#include "bzbutil/platform.h"
#include "bzbstring/stringutils.h"
#include "bzbfile/elfparser.h"

#if defined(WINDOWS)
#define WINDOWS_LEAN_AND_MEAN 1

#include <windows.h>
#include <tlhelp32.h>

using IsWow64Process_ptr = BOOL(WINAPI*)(HANDLE, PBOOL);

#elif defined(LINUX)

#include <unistd.h>
#include <sys/uio.h>
#include <sys/mman.h>
#include <linux/limits.h>
#include <sys/types.h>
#include <signal.h>

#endif


int bzbmem::get_pid() noexcept {

#if defined(WINDOWS)

	return GetCurrentProcessId();

#elif defined(LINUX)

	return getpid();

#endif

}


std::vector<int> bzbmem::get_pid_by_name(const std::string &name) {

	std::vector<int> output;

#if defined(WINDOWS)

	PROCESSENTRY32 processEntry;
	HANDLE handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

	if (!handle) {
		bzbutil::throw_sys_error();
	}

	// We must initialize this word, otherwise Process32First fails.
	processEntry.dwSize = sizeof(PROCESSENTRY32);

	if (Process32First(handle, &processEntry)) {
		do {
			if (!lstrcmp(processEntry.szExeFile, name.c_str()))
				output.push_back(processEntry.th32ProcessID);
		} while (Process32Next(handle, &processEntry));
	} else {
		bzbutil::throw_sys_error();
	}

	CloseHandle(handle);

	return output;

#elif defined(LINUX)

	std::vector<std::string> procs = bzbfile::get_dir_folder_list("/proc/");

	char comm_path[PATH_MAX];
	for (const auto &proc_dir : procs) {

		if (std::all_of(proc_dir.begin(), proc_dir.end(),
			[](char c){ return (c >= '0' && c <= '9'); })) {

			snprintf(comm_path, PATH_MAX, "/proc/%s/comm", proc_dir.c_str());
			std::string content = bzbfile::read(comm_path);
			bzbstring::chomp(content);

			if (content == name) {
				output.push_back(std::stoi(proc_dir));
			}
		}

	}

	return output;

#endif

}


void bzbmem::kill(int pid) {

#if defined(WINDOWS)

	HANDLE proc_handle = OpenProcess(PROCESS_TERMINATE, false, pid);
	int res = TerminateProcess(proc_handle, 0);
	if(!res) {
		bzbutil::throw_sys_error();
	}

#elif defined(LINUX)

	int res = ::kill(pid, SIGTERM);
	if(res) {
		bzbutil::throw_sys_error();
	}

#endif

}


void bzbmem::sleep(unsigned int millis) {

#ifdef WINDOWS

	Sleep(millis);

#elif defined(LINUX)

	usleep(millis * 1000);

#endif

}


std::vector<char> bzbmem::read(int pid, size_t address, size_t size) {

	std::vector<char> output;

#if defined(WINDOWS)

	HANDLE process_handle = OpenProcess(PROCESS_VM_READ, FALSE, pid);
	if(!process_handle) {
		bzbutil::throw_sys_error();
	}

	output.resize(size);

	if (!ReadProcessMemory(process_handle, (LPVOID) address,
			&output[0], size, (SIZE_T*) nullptr)) {

		bzbutil::throw_sys_error();
	}

	CloseHandle(process_handle);

	return output;

#elif defined(LINUX)

	iovec local;
	iovec remote;

	local.iov_base = &output[0];
	local.iov_len = size;

	remote.iov_base = (void*) address;
	remote.iov_len = size;

	ssize_t bytes = process_vm_readv(pid, &local, 1, &remote, 1, 0);
	if (bytes < 0) {
		bzbutil::throw_sys_error();
	}

	return output;

#endif

}


size_t bzbmem::write(int pid, size_t address, std::vector<char> &data, size_t size) {

#if defined(WINDOWS)

	HANDLE process_handle;
	SIZE_T bytecount = 0;
	LPCVOID value_ptr = static_cast<LPCVOID>(&data[0]);

	process_handle = OpenProcess(PROCESS_ALL_ACCESS, FALSE, pid);
	if(!process_handle) {
		bzbutil::throw_sys_error();
	}

	if (!WriteProcessMemory(process_handle, (LPVOID) address,
			value_ptr, static_cast<SIZE_T>(size), &bytecount)) {

		bzbutil::throw_sys_error();
	}

	CloseHandle(process_handle);

	return bytecount;

#elif defined(LINUX)

	iovec local;
	iovec remote;

	local.iov_base = &data[0];
	local.iov_len = size;

	remote.iov_base = (void*) address;
	remote.iov_len = size;

	ssize_t bytes = process_vm_writev(pid, &local, 1, &remote, 1, 0);
	if (bytes < 0) {
		bzbutil::throw_sys_error();
	}

	return (int) bytes;

#endif
}


void bzbmem::lock(void *address, size_t size) {

#if defined(WINDOWS)

	int res = VirtualLock(address, size);
	if(!res) {
		bzbutil::throw_sys_error();
	}

#elif defined(LINUX)

	int res = mlock(address, size);
	if(res) {
		bzbutil::throw_sys_error();
	}

#endif

}


void bzbmem::unlock(void *address, size_t size) {

#if defined(WINDOWS)

	int res = VirtualUnlock(address, size);
	if(!res) {
		bzbutil::throw_sys_error();
	}

#elif defined(LINUX)

	int res = munlock(address, size);
	if(res) {
		bzbutil::throw_sys_error();
	}

#endif

}


#ifdef WINDOWS

int bzbmem::get_pid_by_win_name(const std::string& win_name) {

	HWND win_handle = FindWindow(nullptr, win_name.c_str());
	if(!win_handle) {
		return -1;
	}

	DWORD process_id;
	GetWindowThreadProcessId(win_handle, &process_id);

	return process_id;
}


int bzbmem::inject_dll(int pid, std::string dllname) {

	if(pid < 1) {
		bzbutil::throw_sys_error(ERROR_INVALID_DATA);
	}

	HANDLE process = OpenProcess(PROCESS_CREATE_THREAD | PROCESS_QUERY_INFORMATION |
								PROCESS_VM_OPERATION | PROCESS_VM_WRITE | PROCESS_VM_READ,
								false, pid);

	if(!process) {
		bzbutil::throw_sys_error();
	}

	// LoadLibraryA will be most likely already loaded and mapped to the right address
	void *LoadLibraryA_ptr = (void*) GetProcAddress(GetModuleHandle("kernel32.dll"), "LoadLibraryA");

	if(!LoadLibraryA_ptr) {
		bzbutil::throw_sys_error();
	}

	// Allocate space for the DLL's name on the target process
	void *remote_str = (void*) VirtualAllocEx(process, nullptr,
												dllname.size(),
												MEM_RESERVE | MEM_COMMIT,
												PAGE_EXECUTE_READWRITE);
	if(!remote_str) {
		bzbutil::throw_sys_error();
	}

	// Write the string to the allocated space
	bool write_ret = WriteProcessMemory(process, (void*) remote_str,
										dllname.c_str(), dllname.size(), nullptr);

	if(!write_ret) {
		bzbutil::throw_sys_error();
	}

	int new_pid;
	// Start the thread passing LoadLibraryA as entry function
	// LoadLibraryA will use the DLL's name as parameter
	HANDLE thread = CreateRemoteThread(process, nullptr, 0,
										(LPTHREAD_START_ROUTINE) LoadLibraryA_ptr,
										remote_str, 0, (DWORD*) &new_pid);

	if(!thread) {
		bzbutil::throw_sys_error();
	}

	CloseHandle(process);
	CloseHandle(thread);
	return new_pid;

}


int bzbmem::inject_code(int pid, void* code, size_t size, void* args) {

	if(pid < 1) {
		bzbutil::throw_sys_error(ERROR_INVALID_DATA);
	}

	HANDLE process = OpenProcess(PROCESS_CREATE_THREAD | PROCESS_QUERY_INFORMATION |
								PROCESS_VM_OPERATION | PROCESS_VM_WRITE | PROCESS_VM_READ,
								false, pid);

	if(!process) {
		bzbutil::throw_sys_error();
	}

	// Allocate space for the code on the target process
	void *remote_str = (void*) VirtualAllocEx(process, nullptr, size,
												MEM_RESERVE | MEM_COMMIT,
												PAGE_EXECUTE_READWRITE);
	if(!remote_str) {
		bzbutil::throw_sys_error();
	}

	// Write the string to the allocated space
	bool write_ret = WriteProcessMemory(process, (void*) remote_str,
										code, size, nullptr);

	if(!write_ret) {
		bzbutil::throw_sys_error();
	}

	int new_pid;
	HANDLE thread = CreateRemoteThread(process, nullptr, 0,
										(LPTHREAD_START_ROUTINE) remote_str,
										args, 0, (DWORD*) &new_pid);

	if(!thread) {
		bzbutil::throw_sys_error();
	}

	CloseHandle(process);
	CloseHandle(thread);
	return new_pid;
}

#endif


bool bzbmem::is_64bit_proc(int pid) {

#ifdef WINDOWS

	BOOL is64bit = FALSE;
	HANDLE process = OpenProcess(PROCESS_QUERY_INFORMATION, false, pid);

	IsWow64Process_ptr IsWow64Process = (IsWow64Process_ptr) GetProcAddress(
		GetModuleHandle(TEXT("kernel32")),"IsWow64Process");

	if(IsWow64Process) {
		if (!IsWow64Process(process, &is64bit)) {
			bzbutil::throw_sys_error();
		} else {

			if(is64bit) {
				return false;
			} else {
				SYSTEM_INFO info;
				GetNativeSystemInfo(&info);
				if(info.wProcessorArchitecture == PROCESSOR_ARCHITECTURE_AMD64) {
					return true;
				} else {
					return false;
				}
			}
		}
	} else {
		return false;
	}

#elif defined(LINUX)

	std::string filename = bzbfile::readlink(
							bzbstring::concat("/proc/", pid, "/exe"));

	if(bzbfile::get_elf_class(filename) == bzbfile::ELF_CLASS_64)
		return true;

	return false;

#endif

}
